package trainer;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Random;
/**
 * Interpolater
 * @author Yunkai Zhang
 *
 */
public class Interpolater {

  private String[] types = { "Unigram", "Bigram", "Trigram", "Interpolated Model" };

  double lprev = Double.NEGATIVE_INFINITY;

  double lcur = Double.NEGATIVE_INFINITY;

  double ltri = 0;

  double lbi = 0;

  double luni = 0;

  int size = 0;

  private int N = 3;

  private double E = 0.00001;

  /**
   * Constructor
   * 
   * @param N
   * @param E
   */
  public Interpolater(int N, double E) {
    this.N = N;
    this.E = E;
  }

  /**
   * Use EM to generate the weights
   * 
   * @param models
   * @param heldoutdata
   * @param evaldata
   * @param winnerCounts
   * @return
   */
  public double[] Interpolate(ArrayList<HashMap<String, Double>> models,
          ArrayList<String> heldoutdata, ArrayList<String> evaldata, int[] winnerCounts) {

    double[] tmpWeight = initialWeight(N);

    int count = 1;

    lcur = Double.NEGATIVE_INFINITY;

    do {
      count++;
      lprev = lcur;

      tmpWeight = updateWeight(tmpWeight, models, heldoutdata);

    } while (!isCoverge(lprev, lcur));
    System.out.println("Iteration: " + count);

    System.out.print("PP on the dataset for EM");
    System.out.print(" " + Math.exp(-(luni / size)) + " " + Math.exp(-(lbi / size)) + " "
            + Math.exp(-(ltri / size)) + " ");

    System.out.println(" " + Math.exp(-lcur) + " " + tmpWeight[0] + " " + tmpWeight[1] + " "
            + tmpWeight[2]);

    int bestModel = compareModel(models, tmpWeight, evaldata);
    winnerCounts[bestModel]++;
    System.out.println("Winner is: " + types[bestModel]);

    return tmpWeight;

  }

  /**
   * Return the index of the best model with larges avg-ll
   * 
   * @param ll
   * @return
   */
  private int max(double[] ll) {
    int index = 0;
    double maxValue = Double.NEGATIVE_INFINITY;
    for (int i = 0; i < ll.length; i++) {
      if (maxValue < ll[i]) {
        index = i;
        maxValue = ll[i];
      }
    }
    return index;
  }

  /**
   * Randomly select weights and normalize them to sum to 1.
   * 
   * @param N
   * @return
   */
  private double[] initialWeight(int N) {
    double[] weight = new double[N];
    double sum = 0;
    Random r = new Random();
    for (int i = 0; i < N; i++) {
      weight[i] = r.nextDouble();
      sum += weight[i];
    }

    for (int i = 0; i < N; i++) {
      weight[i] = weight[i] / sum;
    }

    return weight;
  }

  /**
   * To judge whether the EM process is converged
   * 
   * @param lprev
   * @param lcur
   * @return
   */
  private boolean isCoverge(double lprev, double lcur) {
    double dist = (lcur - lprev) / Math.abs(lcur);
    if (dist < E)
      return true;
    else
      return false;
  }

  /**
   * Updating the weights based on the previous weights.
   * 
   * @param tmpWeight
   * @param models
   * @param heldoutdata
   * @return
   */
  private double[] updateWeight(double[] tmpWeight, ArrayList<HashMap<String, Double>> models,
          ArrayList<String> heldoutdata) {

    HashMap<String, Double> unigram = models.get(0);
    HashMap<String, Double> bigram = models.get(1);
    HashMap<String, Double> trigram = models.get(2);

    double[] weight = new double[N];

    double unigramScore;
    double bigramScore;
    double trigramScore;
    double scoreSum;
    String cur;
    String w1;
    String w2;

    lcur = 0;

    ltri = 0;
    lbi = 0;
    luni = 0;

    size = 0;

    for (int i = 2; i < heldoutdata.size(); i++) {
      cur = heldoutdata.get(i);

      if (cur.equalsIgnoreCase("<S>"))
        continue;

      w1 = heldoutdata.get(i - 2);
      w2 = heldoutdata.get(i - 1);

      size++;

      Double unigramProb = unigram.get(cur);
      if (unigramProb == null) {
        unigramProb = 0.0;
      }

      Double bigramProb = bigram.get(w2 + " " + cur);
      if (bigramProb == null) {
        bigramProb = 0.0;
      }

      Double trigramProb = trigram.get(w1 + " " + w2 + " " + cur);
      if (trigramProb == null) {
        trigramProb = 0.0;
      }

      unigramScore = unigramProb * tmpWeight[0];
      bigramScore = bigramProb * tmpWeight[1];
      trigramScore = trigramProb * tmpWeight[2];

      scoreSum = unigramScore + bigramScore + trigramScore;

      weight[0] += unigramScore / scoreSum;
      weight[1] += bigramScore / scoreSum;
      weight[2] += trigramScore / scoreSum;

      if (unigramScore != 0)
        luni += Math.log(unigramScore / tmpWeight[0]);
      if (bigramScore != 0)
        lbi += Math.log(bigramScore / tmpWeight[1]);
      if (trigramScore != 0)
        ltri += Math.log(trigramScore / tmpWeight[2]);
      else {
        System.out.println(w1 + " " + w2 + " " + cur);
      }

      lcur += Math.log(scoreSum);
    }

    lcur = lcur / size;

    weight[0] = weight[0] / size;
    weight[1] = weight[1] / size;
    weight[2] = weight[2] / size;

    return weight;
  }

  /**
   * Evaluate the model by using the evaldata set, which is not seen by the all the models and
   * parameters
   * 
   * @param models
   * @param tmpWeight
   * @param evaldata
   * @return
   */
  private int compareModel(ArrayList<HashMap<String, Double>> models, double[] tmpWeight,
          ArrayList<String> evaldata) {

    HashMap<String, Double> unigram = models.get(0);
    HashMap<String, Double> bigram = models.get(1);
    HashMap<String, Double> trigram = models.get(2);

    double ltri = 0;
    double lbi = 0;
    double luni = 0;

    ArrayList<String> allData = evaldata;

    double lcur = 0;

    int size = 0;

    for (int i = 2; i < allData.size(); i++) {
      String cur = allData.get(i);

      if (cur.equalsIgnoreCase("<S>"))
        continue;
      String w1 = allData.get(i - 2);
      String w2 = allData.get(i - 1);
      size++;

      Double unigramProb = unigram.get(cur);
      if (unigramProb == null) {
        unigramProb = 0.0;
      }

      Double bigramProb = bigram.get(w2 + " " + cur);
      if (bigramProb == null) {
        bigramProb = 0.0;
      }

      Double trigramProb = trigram.get(w1 + " " + w2 + " " + cur);
      if (trigramProb == null) {
        trigramProb = 0.0;
      }

      double unigramScore = unigramProb * tmpWeight[0];
      double bigramScore = bigramProb * tmpWeight[1];
      double trigramScore = trigramProb * tmpWeight[2];

      luni += Math.log(unigramProb);
      lbi += Math.log(bigramProb);
      ltri += Math.log(trigramProb);
      lcur += Math.log(unigramScore + bigramScore + trigramScore);

    }

    lcur = lcur / size;
    luni = luni / size;
    lbi = lbi / size;
    ltri = ltri / size;

    System.out.print("PP on evaluation folder");
    System.out.print(" " + Math.exp(-(luni)) + " " + Math.exp(-(lbi)));
    System.out.print(" " + Math.exp(-(ltri)) + " ");
    System.out.println(" " + Math.exp(-lcur));
    double[] avgll = { luni, lbi, ltri, lcur };

    return max(avgll);

  }
}
